package common

import (
	"sync"
	"time"
)

// SecurityConfig 安全配置
type SecurityConfig struct {
	// IP限流配置
	MaxRequestsPerSecond int           // 每秒最大请求数
	BlockDuration        time.Duration // IP被封禁的时长
	MaxConnPerIP         int           // 每个IP的最大连接数

	// 请求大小限制
	MaxRequestSize int64 // 最大请求体积(bytes)

	// 黑白名单
	WhiteList []string // IP白名单
	BlackList []string // IP黑名单
}

// SecurityManager 安全管理器
type SecurityManager struct {
	config *SecurityConfig

	// IP访问计数器
	ipCounter     map[string]int
	ipCounterLock sync.RWMutex

	// IP封禁列表
	bannedIPs     map[string]time.Time
	bannedIPsLock sync.RWMutex

	// 连接计数器
	connCounter map[string]int
	connLock    sync.RWMutex
}

// NewSecurityManager 创建安全管理器
func NewSecurityManager(config *SecurityConfig) *SecurityManager {
	sm := &SecurityManager{
		config:      config,
		ipCounter:   make(map[string]int),
		bannedIPs:   make(map[string]time.Time),
		connCounter: make(map[string]int),
	}

	// 启动清理计时器
	go sm.cleanupLoop()

	return sm
}

// IsIPAllowed 检查IP是否允许访问
func (sm *SecurityManager) IsIPAllowed(ip string) bool {
	// 检查白名单
	for _, whiteIP := range sm.config.WhiteList {
		if ip == whiteIP {
			return true
		}
	}

	// 检查黑名单
	for _, blackIP := range sm.config.BlackList {
		if ip == blackIP {
			return false
		}
	}

	// 检查是否被封禁
	sm.bannedIPsLock.RLock()
	if banTime, ok := sm.bannedIPs[ip]; ok {
		if time.Now().Before(banTime) {
			sm.bannedIPsLock.RUnlock()
			return false
		}
		// 解封
		sm.bannedIPsLock.RUnlock()
		sm.bannedIPsLock.Lock()
		delete(sm.bannedIPs, ip)
		sm.bannedIPsLock.Unlock()
	} else {
		sm.bannedIPsLock.RUnlock()
	}

	return true
}

// IncrementRequestCount 增加请求计数
func (sm *SecurityManager) IncrementRequestCount(ip string) bool {
	sm.ipCounterLock.Lock()
	defer sm.ipCounterLock.Unlock()

	count := sm.ipCounter[ip]
	if count >= sm.config.MaxRequestsPerSecond {
		// 触发封禁
		sm.banIP(ip)
		return false
	}

	sm.ipCounter[ip]++
	return true
}

// IncrementConnCount 增加连接计数
func (sm *SecurityManager) IncrementConnCount(ip string) bool {
	sm.connLock.Lock()
	defer sm.connLock.Unlock()

	count := sm.connCounter[ip]
	if count >= sm.config.MaxConnPerIP {
		return false
	}

	sm.connCounter[ip]++
	return true
}

// DecrementConnCount 减少连接计数
func (sm *SecurityManager) DecrementConnCount(ip string) {
	sm.connLock.Lock()
	defer sm.connLock.Unlock()

	if sm.connCounter[ip] > 0 {
		sm.connCounter[ip]--
	}
}

// banIP 封禁IP
func (sm *SecurityManager) banIP(ip string) {
	sm.bannedIPsLock.Lock()
	defer sm.bannedIPsLock.Unlock()

	sm.bannedIPs[ip] = time.Now().Add(sm.config.BlockDuration)
}

// cleanupLoop 定期清理计数器
func (sm *SecurityManager) cleanupLoop() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for range ticker.C {
		// 清理请求计数器
		sm.ipCounterLock.Lock()
		sm.ipCounter = make(map[string]int)
		sm.ipCounterLock.Unlock()

		// 清理过期的封禁IP
		now := time.Now()
		sm.bannedIPsLock.Lock()
		for ip, banTime := range sm.bannedIPs {
			if now.After(banTime) {
				delete(sm.bannedIPs, ip)
			}
		}
		sm.bannedIPsLock.Unlock()
	}
}
